from classify.bayes.basic import *
from classify.bayes.continuity.nb_functions import NBFunctions


class GaussianNB(NaiveBayes):
    def feed_data(self, x, y, sample_weight=None):

        x = np.array([list(map(lambda c: float(c), sample)) for sample in x])

        labels = list(set(y))
        label_dic = {label: i for i, label in enumerate(labels)}
        y = np.array([label_dic[yy] for yy in y])
        cat_counter = np.bincount(y)
        labels = [y == value for value in range(len(cat_counter))]
        labelled_x = [x[label].T for label in labels]

        self._x, self.y = x.T, y
        self._labelled_x, self._label_zip = labelled_x, labels
        self._cat_counter, self.label_dic = cat_counter, {i: _l for _l, i in label_dic.items()}
        self.feed_sample_weight(sample_weight)

    def feed_sample_weight(self, sample_weight=None):
        if sample_weight is not None:
            local_weights = sample_weight * len(sample_weight)
            for i, label in enumerate(self._label_zip):
                self._labelled_x[i] *= local_weights[label]

    def _fit(self, lb):
        n_category = len(self._cat_counter)
        p_category = self.get_prior_probability(lb)

        data = [
            NBFunctions.gaussian_maximum_likelihood(
                self._labelled_x, n_category, dim) for dim in range(len(self._x))]
        self._data = data
        def func(input_x, tar_category):
            rs = 1
            for d, xx in enumerate(input_x):

                rs *= data[d][tar_category](xx)
            return rs * p_category[tar_category]
        return func

    @staticmethod
    def _transfer_x(x):
        return x